# -*- coding: utf-8 -*-
"""
# -*- coding: utf-8 -*-
Created on Sat Feb 20 18:42:13 2021
this code is modified from mnist_cnn for detecting cat and dog
@author: LI
"""

import tensorflow as tf
import matplotlib.pyplot as plt
import glob
import time
import sys
import numpy as np
from skimage import io, transform
#from tensorflow.python.framework import graph_util
import cv2 as cv
import os
import os.path
import copy

import stat
import shutil

import Alaxnet_CNN
import LoadData
import config

row = config.row
col = config.col
c = config.c
class_num = config.class_num

def _show_time(cost_time):
    #start_time = time.time()
    hour =   ( cost_time // 60 ) // 60
    minute = ( (cost_time) // 60 ) % 60
    second = (  cost_time) % 60

    return hour, minute, second

def delete_file(filePath):
    if os.path.exists(filePath):
        for fileList in os.walk(filePath):
            for name in fileList[2]:
                os.chmod(os.path.join(fileList[0], name), stat.S_IWRITE)
                os.remove(os.path.join(fileList[0], name))
        shutil.rmtree(filePath)
        return "delete ok"
    else:
        return "no filepath"

# 定义两个placeholder

x = tf.placeholder(tf.float32, shape=[None, row, col, c], name='x-input')
y = tf.placeholder(tf.int64, shape=[None, class_num], name='y-input')
 # keep_prob用来表示神经元的输出概率
keep_prob = tf.placeholder(tf.float32, name='keep_prob')

cost = []
# 设置批次的大小
batch_size = config.batch_size
# 计算一共有多少个批次
epochs = config.epochs
learnRate = config.learnRate
keepPro_train = config.keepPro_train


# 载入数据集path
path = config.path
save_path = config.save_path
# 创建会话
with tf.Session() as sess:
    model = Alaxnet_CNN.ALAXNET(x, y, keep_prob, learnRate, class_num)
# >>>   read img data 
    data, label = LoadData.read_img(path)
    num_example = data.shape[0]
    arr = np.arange(num_example)  
    np.random.shuffle(arr)
    data = data[arr]
    label = label[arr]  # 打乱顺序
# <<<

    exit(0)
    plt.ion()
    plt.figure(1)
    #start_time = time.clock()
    start_time = time.time()
    sess.run(tf.global_variables_initializer())  # 初始化变量   
    for epoch in range(epochs):  # 迭代epochs次（训练21次）
        for batch in range(0, len(data), batch_size):
            batch_xs = data[batch:batch + batch_size]
            batch_ys = label[batch:batch + batch_size]
            batch_xs = np.array(batch_xs)
            sess.run(model.train_step, feed_dict={x: batch_xs, y: batch_ys, keep_prob: keepPro_train})  # 进行迭代训练
   
        print("迭代： ", epoch+1, " 次")
        cost_value = sess.run(model.cross_entropy, feed_dict={x: batch_xs, y: batch_ys, keep_prob: keepPro_train})
        train_acc, pre = sess.run([model.accuracy, model.prediction], feed_dict={x: batch_xs, y: batch_ys, keep_prob: keepPro_train})
        cost.append(cost_value)
        print(pre.shape, pre[0][0])
        print('loss cost: ', cost_value)
        print('Training Accuracy=           '+str(train_acc))
        # 绘制损失曲线
        plt.clf()
        plt.plot(cost, 'g-')
        plt.pause(0.01)

        if (train_acc > 0.9) and (epoch % 1000 == 0):
        #save model cpkt
            print(delete_file(save_path))  # detect the file_dir, delete it if it exits
            ckpt_file_path = save_path
            path_ = os.path.dirname(os.path.abspath(ckpt_file_path))
            if os.path.isdir(path_) is False:
                os.makedirs(path_)
            saver = tf.train.Saver(max_to_keep=3)
            saver.save(sess, ckpt_file_path + '/model.ckpt', write_meta_graph=True)
            
        #save model pb
            graph = tf.compat.v1.graph_util.convert_variables_to_constants(sess, sess.graph_def, ["x-input","y-input","softmax","keep_prob"])
            tf.io.write_graph(graph, './models', 'model.pb',as_text= False)
            plt.plot(cost, 'g-')
            plt.savefig("./loss.jpg")

    cost_time = time.time() - start_time
    
    hour =   ( cost_time // 60 ) // 60
    minute = ( (cost_time) // 60 ) % 60
    second = (  cost_time) % 60
    
    print('Running time:%f Second' % cost_time)  # 输出运行时间
    print('Running time: {:.0f}h {:.0f}m {:.0f}s '.format(hour, minute, second))
    sys.exit(0)
#>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>


    # 测试数据计算出准确率
    print("")

    for i in range(2):
        image_idx = i + 1
        # train
        test_img, test_label, test_show = test_readImg(path, image_idx)
        # test
        #test_img,test_label,test_show = test_readImg(path_test,image_idx)

        prediction = sess.run(model.prediction, feed_dict={
                              x: test_img, y: test_label, keep_prob: 1.0})

        pre_idx = np.argmax(prediction)

        print('real:    ', test_label)
        print('Training prediction=', prediction)
        print(' -> output class is: ', pre_idx + 1, "\n",
              '-> prediction   is: ', prediction[0][pre_idx])
        print('>>>>>>>>>>>>>>>\t end \t<<<<<<<<<<<<<')
        cv.imshow('aa', test_show)
        cv.waitKey(0)
        cv.destroyAllWindows()
# 0  4   2   3   1
